# -*- coding:utf8 -*-

import os
import random
# from PIL import Image

allowed_names = ['jpg', 'JPG']  # 需要随机替换的后缀名列表
COUNT_NUMBER = 0
TRAIN_NUMBER_RATE = 0.8
SAVE_PATH = '/home/chenzl/my/objection_detection_under_water/VOC2020/ImageSets/Main/'
INPUT_PATH = '/home/chenzl/my/objection_detection_under_water/VOC2020/Annotation/'
TRAIN_TXT_NAME = 'train.txt'
VAL_TXT_NAME = 'val.txt'
DATA_AUGMENTATION = 0

def file_name(name):
    a = name.split('.', 1)[0]
    return name.split('.', 1)[0]

def file_constract(path):
    files = os.listdir(path)  # 获取当前目录的所有文件及文件夹
    file_number = len(files)
    train_number = int(file_number * TRAIN_NUMBER_RATE)
    train_files = random.sample(files, train_number)

    train_file_names = [file_name(train_file) for train_file in train_files]
    # for train_file in train_files:
    #     b = file_name(train_file)
    #     c = 1
    # files.remove(train_file for train_file in train_files)
    for train_file in train_files:
        files.remove(train_file)
    val_file_names = [file_name(val_file) for val_file in files]

    return train_file_names, val_file_names

if __name__ == '__main__':
    # os.system('cp -R /home/user/项目代码文件/Annotation /usr/VOCdevkit/VOC2019/')
    os.system('rm /home/chenzl/my/objection_detection_under_water/VOC2020/ImageSets/Main/train.txt')
    os.system('rm /home/chenzl/my/objection_detection_under_water/VOC2020/ImageSets/Main/val.txt')
    train_names, val_names = file_constract(INPUT_PATH)
    with open(SAVE_PATH+TRAIN_TXT_NAME, 'a') as f:
        for train_name in train_names:
            f.write(str(train_name)+'\n')
            if DATA_AUGMENTATION == 1:
                f.write(str(train_name) + 'LR90'+'\n')
                f.write(str(train_name) + 'RR90'+'\n')
                f.write(str(train_name) + 'CTB'+'\n')
                f.write(str(train_name) + 'CLR'+'\n')
        f.close()
    with open(SAVE_PATH+VAL_TXT_NAME, 'a') as f:
        for val_name in val_names:
            f.write(str(val_name)+'\n')
            if DATA_AUGMENTATION == 1:
                f.write(str(val_name) + 'LR90'+'\n')
                f.write(str(val_name) + 'RR90'+'\n')
                f.write(str(val_name) + 'CTB'+'\n')
                f.write(str(val_name) + 'CLR'+'\n')
        f.close()
    if DATA_AUGMENTATION == 1:
        os.system('python voc_label2.py --DA True')
    else:
        os.system('python voc_label2.py')